# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: chensili

import functools
import json
import os
import traceback

import MySQLdb
from tabulate import tabulate

from coin.base.timestamp import get_timestamp
from coin.base.book.book_builder_base import BookBuilderBase
from coin.base.book.types import MarketTradeEvent
from coin.support.feed_tool.feed_stats.logic.util import *
from coin.support.feed_tool.feed_stats.logic.book_stats_calculator import (gen_empty_BookStats,
                                                                           BookStatsCalculator)
from coin.support.feed_tool.feed_stats.logic.trade_stats_calculator import (gen_empty_TradeStats,
                                                                            TradeStatsCalculator)
from coin.support.feed_tool.feed_stats.logic.cross_stats_calculator import (gen_empty_CrossStats,
                                                                            CrossStatsCalculator)


class DumpStrategy(object):
  def __init__(self,
               machine,
               products,
               trading_date: datetime.datetime,
               start_time: datetime.datetime,
               end_time: datetime.datetime,
               git_commit_datetime: datetime.datetime,
               git_commit_sha: str,
               sub_request):
    time_fmt = '%Y-%m-%d %H:%M:%S.%f'
    self._machine = machine
    self._products = products
    self._trading_date = trading_date.strftime('%Y-%m-%d')
    self._start_time = start_time.strftime(time_fmt)
    self._end_time = end_time.strftime(time_fmt)
    self._git_commit_datetime = git_commit_datetime.strftime(time_fmt)
    self._git_commit_sha = git_commit_sha
    self._sub_request = sub_request
    exchange_api_id = get_exchange_api_id_from_request([sub_request])[0]
    self._exchange_api_id = exchange_api_id
    self._book_stats = {}
    self._trade_stats = {}
    self._cross_stats = {}

  def _get_book_stats(self, sub_key, symbol):
    return self._book_stats.get((sub_key, symbol), None)

  def _get_trade_stats(self, sub_key, symbol):
    return self._trade_stats.get((sub_key, symbol), None)

  def _get_cross_stats(self, sub_key, symbol):
    return self._cross_stats.get((sub_key, symbol), None)

  def on_feed_reset(self, sub_key, book_builder):
    feed_checker = book_builder._feed_handler.feed_checker
    for product in self._products:
      if sub_key == self._sub_request:
        feed_handler = functools.partial(self.on_feed, sub_key, feed_checker, product)
        book_builder.subscribe(product, feed_handler)
        book_builder.subscribe_trade(product, feed_handler)
        exchange = sub_key.exchange
        symbol = to_book_builder_symbol(product, self._sub_request)

        book_stats = BookStatsCalculator(gen_empty_BookStats())
        trade_stats = TradeStatsCalculator(gen_empty_TradeStats())
        cross_stats = CrossStatsCalculator(gen_empty_CrossStats())
        book_stats.set_product(product)
        trade_stats.set_product(product)
        cross_stats.set_product(product)
        self._book_stats[(sub_key, symbol)] = book_stats
        self._trade_stats[(sub_key, symbol)] = trade_stats
        self._cross_stats[(sub_key, symbol)] = cross_stats
    feed_checker.start(timestamp=get_timestamp())

  def on_feed(self, sub_key, feed_checker, product, feed):
    exchange = sub_key.exchange
    book_stats = self._get_book_stats(sub_key, feed.symbol)
    trade_stats = self._get_trade_stats(sub_key, feed.symbol)
    cross_stats = self._get_cross_stats(sub_key, feed.symbol)
    assert book_stats is not None, (sub_key, feed.symbol)
    assert trade_stats is not None, (sub_key, feed.symbol)
    assert cross_stats is not None, (sub_key, feed.symbol)
    if isinstance(feed, BookBuilderBase):
      book_stats.update_by_book(feed, feed_checker, product)
      cross_stats.update_by_book(feed)
    elif isinstance(feed, MarketTradeEvent):
      trade_stats.update_by_trade(feed, feed_checker, product)
      cross_stats.update_by_trade(feed)

  def print_book_stats(self):
    stats_table = []
    for (exchange, symbol), calculator in sorted(self._book_stats.items()):
      stats = calculator.book_stats
      start_time = to_datetime_or_none(stats.start_time)
      end_time = to_datetime_or_none(stats.end_time)
      stats_table.append([
          self._machine,
          exchange,
          symbol,
          start_time,
          end_time,
          stats.book_count,
          stats.each_hour_count,
          stats.crossed_book_count,
          stats.mpc_count,
          avg_or_none(stats.bid0_qty_total, stats.book_count),
          avg_or_none(stats.ask0_qty_total, stats.book_count),
          avg_or_none(stats.min_side_qty_total, stats.book_count),
          avg_or_none(stats.max_side_qty_total, stats.book_count),
          avg_or_none(stats.bidask_spread_total, stats.book_count),
          to_datetime_or_none(stats.longest_no_feed.begin),
          stats.longest_no_feed.time_delta / (10.**9),
          stats.feed_checker_invalid.feed_checker_invalid_count,
          stats.feed_checker_invalid.max_feed_checker_consecutive_invalid_count,
          stats.feed_checker_invalid.max_feed_checker_consecutive_invalid_period / (10.**9)
      ])
    if all(table_item[2] is None for table_item in stats_table):
      print('No book feed for this run!')
    else:
      print(
          tabulate(stats_table,
                   headers=[
                       'machine',
                       'exchange',
                       'native_symbol',
                       'start_time',
                       'end_time',
                       'book_count',
                       'each_hour_count',
                       'crossed_book_count',
                       'mpc_count',
                       'bid0_qty_avg',
                       'ask0_qty_avg',
                       'min_side_qty_avg',
                       'max_side_qty_avg',
                       'bidask_spread_avg',
                       'longest_no_feed_begin',
                       'no_feed_seconds',
                       'feed_checker_invalid_count',
                       'max_feed_checker_consecutive_invalid_count',
                       'max_feed_checker_consecutive_invalid_seconds'
                   ]))

  def print_trade_stats(self):
    stats_table = []
    for key, value in sorted(self._trade_stats.items()):
      trade_stats = value.trade_stats
      first_trade = trade_stats.basic_stats.first_trade
      last_trade = trade_stats.basic_stats.last_trade
      start_time = datetime.datetime.fromtimestamp(first_trade /
                                                   (10.**9)) if first_trade else first_trade
      end_time = datetime.datetime.fromtimestamp(last_trade /
                                                 (10.**9)) if last_trade else last_trade
      no_trade_delta = trade_stats.max_time_no_trade.time_delta
      no_trade_delta = no_trade_delta / (10. ** 9) if no_trade_delta \
                       else no_trade_delta
      stats_table.append([
          self._machine,
          key[0],
          key[1],
          start_time,
          end_time,
          trade_stats.basic_stats.trade_count,
          trade_stats.basic_stats.each_hour_count,
          [trade_stats.max_time_no_trade.begin, trade_stats.max_time_no_trade.end, no_trade_delta],
          trade_stats.basic_stats.feed_checker_invalid_count
      ])
    print(
        tabulate(stats_table,
                 headers=[
                     'machine',
                     'exchange',
                     'native_symbol',
                     'start time',
                     'end time',
                     'trade count',
                     'each hour count',
                     'longest no feed',
                     'feed_checker_invalid_count'
                 ]))

  def print_cross_stats(self):
    stats_table = []
    for key, value in sorted(self._cross_stats.items()):
      cross_stats = value.cross_stats
      stats_table.append([
          self._machine,
          key[0],
          key[1],
          cross_stats.black_count,
          cross_stats.avg_black_size,
          cross_stats.buy_sweep_count,
          cross_stats.avg_buy_sweep_size,
          cross_stats.sell_sweep_count,
          cross_stats.avg_sell_sweep_size,
          cross_stats.normal_count,
          cross_stats.avg_normal_size,
          [
              cross_stats.between_book_trade.max_between_book_trade_count,
              cross_stats.between_book_trade.begin,
              cross_stats.between_book_trade.time_delta,
              avg_or_none(cross_stats.between_book_trade.total_between_book_trade_count,
                          cross_stats.between_book_trade.stats_count)
          ],
          [
              cross_stats.outside_bid_ask_trade.max_outside_bid_ask_trade_count,
              cross_stats.outside_bid_ask_trade.begin,
              cross_stats.outside_bid_ask_trade.time_delta,
              avg_or_none(cross_stats.outside_bid_ask_trade.total_outside_bid_ask_trade_count,
                          cross_stats.outside_bid_ask_trade.stats_count)
          ],
          ratio_or_none(self._trade_stats[key].trade_stats.basic_stats.trade_count,
                        self._book_stats[key].book_stats.book_count)
      ])
    print(
        tabulate(
            stats_table,
            headers=[
                'machine',
                'exchange',
                'native_symbol',
                'black_count',
                'avg_black_size',
                'buy_sweep_count',
                'avg_buy_sweep_size',
                'sell_sweep_count',
                'avg_sell_sweep_size',
                'normal_count',
                'avg_normal_size',
                '[max_between_book_trade_count, begin, time_delta, avg_between_book_trade_count]',
                '[max_outside_bid_ask_trade_count, begin, time_delta, avg_outside_bid_ask_trade_count]',
                'trade_book_ratio'
            ]))

  def dump_to_db(self, db_config_path):
    db_config_path = os.path.expanduser(db_config_path)
    assert db_config_path and os.path.isfile(db_config_path), \
      'DB config file does not exist: %s' % db_config_path
    with open(db_config_path, 'r') as open_file:
      db_config = json.load(open_file)
    conn = MySQLdb.connect(host=db_config.get('host', None),
                           db=db_config.get('database', None),
                           user=db_config.get('user', None),
                           passwd=db_config.get('password', None))
    cursor = conn.cursor()

    for (exchange, symbol), value in self._book_stats.items():
      try:
        product = value.book_stats.product
        sql, values = self._gen_insert_feed_stats_info_sql(product, symbol)
        cursor.execute(sql, values)
        sql, values = self._gen_query_feed_stats_id_sql(product)
        cursor.execute(sql, values)
        feed_stats_id = cursor.fetchall()[0][0]
        sql, values = self._gen_insert_book_stats_sql(exchange, symbol, feed_stats_id)
        cursor.execute(sql, values)
        sql, values = self._gen_insert_trade_stats_sql(exchange, symbol, feed_stats_id)
        cursor.execute(sql, values)
        sql, values = self._gen_insert_cross_stats_sql(exchange, symbol, feed_stats_id)
        cursor.execute(sql, values)
        conn.commit()
      except Exception as e:
        print('Fail to dump to db.', exchange, symbol, type(e), e)
        print(traceback.format_exc())
        conn.rollback()
        continue
    conn.close()

  def _gen_insert_feed_stats_info_sql(self, product, native_symbol):
    sql = """
INSERT FeedStatsInfo(
    trading_date, calendar_date, start_time, end_time, machine,
    symbol, native_symbol, git_commit_datetime, git_commit_sha, exchange_api_id)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
"""

    values = (self._trading_date,
              self._trading_date,
              self._start_time,
              self._end_time,
              self._machine,
              product.symbol,
              native_symbol,
              self._git_commit_datetime,
              self._git_commit_sha,
              self._exchange_api_id)
    return (sql, values)

  def _gen_query_feed_stats_id_sql(self, product):
    sql = """
SELECT max(feed_stats_id)
FROM FeedStatsInfo
WHERE machine =%s AND exchange_api_id=%s AND symbol=%s AND trading_date=%s AND
      start_time=%s AND end_time=%s
"""
    values = (self._machine,
              self._exchange_api_id,
              product.symbol,
              self._trading_date,
              self._start_time,
              self._end_time)
    return (sql, values)

  def _gen_insert_book_stats_sql(self, exchange, symbol, feed_stats_id):
    sql = """
INSERT BookStats(
    feed_stats_id, time_first_book, time_last_book, book_count, book_hour_count,
    time_no_book_begin, time_no_book_end, no_book_duration, crossed_book_count,
    mpc_count, avg_bid0_qty, avg_ask0_qty, avg_min_side_qty, avg_max_side_qty,
    avg_bid_ask_spread, price_flip_count, avg_price_flip_interval,
    feed_checker_invalid_count, max_feed_checker_consecutive_invalid_count,
    max_feed_checker_consecutive_invalid_period)
VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
"""
    book_stats_cal = self._get_book_stats(exchange, symbol)
    stats = book_stats_cal.book_stats
    no_feed = book_stats_cal.book_stats.longest_no_feed
    values = (feed_stats_id,
              to_time_str_or_none(stats.start_time),
              to_time_str_or_none(stats.end_time),
              stats.book_count,
              json.dumps(stats.each_hour_count),
              to_time_str_or_none(no_feed.begin),
              to_time_str_or_none(no_feed.end),
              no_feed.time_delta / (10.**9),
              stats.crossed_book_count,
              stats.mpc_count,
              avg_or_none(stats.bid0_qty_total, stats.book_count),
              avg_or_none(stats.ask0_qty_total, stats.book_count),
              avg_or_none(stats.min_side_qty_total, stats.book_count),
              avg_or_none(stats.max_side_qty_total, stats.book_count),
              avg_or_none(stats.bidask_spread_total, stats.book_count),
              stats.price_flip_count,
              stats.avg_price_flip_interval / (10.**9),
              stats.feed_checker_invalid.feed_checker_invalid_count,
              stats.feed_checker_invalid.max_feed_checker_consecutive_invalid_count,
              stats.feed_checker_invalid.max_feed_checker_consecutive_invalid_period / (10.**9))
    return (sql, values)

  def _gen_insert_trade_stats_sql(self, exchange, symbol, feed_stats_id):
    trade_stats_cal = self._get_trade_stats(exchange, symbol)
    basic_stats = trade_stats_cal.trade_stats.basic_stats
    no_feed = trade_stats_cal.trade_stats.max_time_no_trade
    vol_stats = trade_stats_cal.trade_stats.volume_stats
    price_stats = trade_stats_cal.trade_stats.price_stats
    sql = """
INSERT TradeStats(
    feed_stats_id, time_first_trade, time_last_trade, trade_count,
    avg_trade_size, trade_hour_count, time_no_trade_begin,
    time_no_trade_end, no_trade_duration, total_volume, buy_volume,
    sell_volume, pct_return, price_open, price_high, price_low,
    price_close, vwap, twap, feed_checker_invalid_count)
VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
"""
    values = (
        feed_stats_id,
        to_time_str_or_none(basic_stats.first_trade),
        to_time_str_or_none(basic_stats.last_trade),
        basic_stats.trade_count,
        basic_stats.avg_trade_size,
        json.dumps(basic_stats.each_hour_count),
        to_time_str_or_none(no_feed.begin),
        to_time_str_or_none(no_feed.end),
        no_feed.time_delta / (10. ** 9) if no_feed.time_delta else \
          no_feed.time_delta,
        vol_stats.total_volume,
        vol_stats.buy_volume,
        vol_stats.sell_volume,
        price_stats.pct_return,
        price_stats.price_open,
        price_stats.price_high,
        price_stats.price_low,
        price_stats.price_close,
        price_stats.vwap,
        price_stats.twap,
        basic_stats.feed_checker_invalid_count)
    return (sql, values)

  def _gen_insert_cross_stats_sql(self, exchange, symbol, feed_stats_id):
    cross_stats_cal = self._get_cross_stats(exchange, symbol)
    stats = cross_stats_cal.cross_stats
    book_count = self._get_book_stats(exchange, symbol).book_stats.book_count
    trade_count = self._get_trade_stats(exchange, symbol).trade_stats.basic_stats.trade_count
    sql = """
INSERT CrossStats(
    feed_stats_id, normal_count, avg_normal_size, black_count, avg_black_size,
    buy_sweep_count, avg_buy_sweep_size, sell_sweep_count, avg_sell_sweep_size,
    max_between_book_trade_count,
    max_between_book_trade_begin,
    max_between_book_trade_duration,
    avg_between_book_trade_count,
    max_outside_bid_ask_trade_count,
    max_outside_bid_ask_trade_begin,
    max_outside_bid_ask_trade_duration,
    avg_outside_bid_ask_trade_count,
    trade_book_ratio)
VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
"""
    values = (feed_stats_id,
              stats.normal_count,
              stats.avg_normal_size,
              stats.black_count,
              stats.avg_black_size,
              stats.buy_sweep_count,
              stats.avg_buy_sweep_size,
              stats.sell_sweep_count,
              stats.avg_sell_sweep_size,
              stats.between_book_trade.max_between_book_trade_count,
              to_time_str_or_none(stats.between_book_trade.begin),
              stats.between_book_trade.time_delta /
              (10.**9) if stats.between_book_trade.time_delta else None,
              avg_or_none(stats.between_book_trade.total_between_book_trade_count,
                          stats.between_book_trade.stats_count),
              stats.outside_bid_ask_trade.max_outside_bid_ask_trade_count,
              to_time_str_or_none(stats.outside_bid_ask_trade.begin),
              stats.outside_bid_ask_trade.time_delta /
              (10.**9) if stats.outside_bid_ask_trade.time_delta else None,
              avg_or_none(stats.outside_bid_ask_trade.total_outside_bid_ask_trade_count,
                          stats.outside_bid_ask_trade.stats_count),
              ratio_or_none(trade_count, book_count))
    return (sql, values)
